from common.variables import get_zmq_url, TimeFrameType
import time
import zmq


def get_kline_topic(symbol: str, timeframe: TimeFrameType) -> str:
    return f"kline_{symbol}_{timeframe}".replace(":", "_").replace("/", "_")


def get_factor_topic(symbol: str, tfname: str, timeframe: TimeFrameType) -> str:
    return f"factor_{symbol}_{tfname}_{timeframe}".replace(":", "_").replace("/", "_")


def get_whole_factor_topic(timeframe: TimeFrameType) -> str:
    return f"whole_factor_{timeframe}"


class send_topic:
    def __init__(self, topic: str):
        self.topic = topic

    def __enter__(self):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUB)
        self.socket.bind(get_zmq_url())
        time.sleep(1)
        return self

    def send(self, msg: str):
        s = f"{self.topic}{msg}"
        print(f"send msg: {time.time()} {s}")
        self.socket.send_string(s)

    def __exit__(self, exc_type, exc_value, traceback):
        self.socket.close()
        self.context.term()


class iter_topic:
    def __init__(self, topic: str):
        self.topic = topic

    def __enter__(self):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.SUB)
        self.socket.connect(get_zmq_url())
        self.socket.setsockopt_string(zmq.SUBSCRIBE, self.topic)
        time.sleep(1)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def __iter__(self):
        return self

    def __next__(self):
        try:
            message = self.socket.recv_string()
            return message[len(self.topic) :]
        except zmq.ZMQError as e:
            raise StopIteration from e

    def close(self):
        self.socket.close()
        self.context.term()
